from langchain_openai import ChatOpenAI
from langchain.agents import tool, Tool, initialize_agent
from langchain.prompts import ChatPromptTemplate
from langchain_core.utils.function_calling import convert_to_openai_function
from pydantic import BaseModel
import json

with open('package.json', 'r') as f:
    config = json.load(f)

model = config['model']
base_url = config['base_url']
api_key = config['api_key']
verbose = bool(config['verbose'])

class CityWeatherInput(BaseModel):
    city: str


@tool(args_schema=CityWeatherInput)
def get_weather(city: str) -> str:
    """Get weather by city"""
    city = city.strip("'")
    weather_data = {
        "beijing": {
            "location": "Beijing",
            "temperature": {
                "current": 32,
                "low": 26,
                "high": 35
            },
            "rain_probability": 10,  # 百分比
            "humidity": 40  # 百分比
        },
        "shenzhen": {
            "location": "Shenzhen",
            "temperature": {
                "current": 28,
                "low": 24,
                "high": 31
            },
            "rain_probability": 90,  # 百分比
            "humidity": 85  # 百分比
        }
    }

    city_key = city.lower()
    if city_key in weather_data:
        return json.dumps(weather_data[city_key], ensure_ascii=False)
    return json.dumps({"error": "Weather Unavailable"}, ensure_ascii=False)


prompt = ChatPromptTemplate.from_messages([
    ("system", "You are helpful by sassy assistant"),
    ("user", "{input}")
])

tools = [get_weather]
functions = [convert_to_openai_function(f) for f in tools]

model = ChatOpenAI(
    model=model,
    base_url=base_url,
    api_key=api_key
)
weather_tool = Tool(name="get_weather", func=get_weather, description="Get weather by city")
agent = initialize_agent([weather_tool], model, verbose=verbose)

result = agent.invoke({"input": "Check the weather in Shenzhen and tell me in one sentence whether I should bring an "
                                "umbrella when going out"})
print(result['output'])
